package ssh

import (
	"bufio"
	"bytes"
	"fmt"
	"git.phjr.com/devops/pkg/util"
	"github.com/pkg/sftp"
	"golang.org/x/crypto/ssh"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path"
	"regexp"
)

type SSHClient struct {
	host        string
	port        string
	username    string
	password    string
	client      *ssh.Client
	sshSession  *ssh.Session
	sftpSession *sftp.Client
}

/**
 * 初始化ssh连接数据
 */
func NewClient(host string, port string, username string, password string) *SSHClient {
	sc := &SSHClient{
		host,
		port,
		username,
		password,
		nil,
		nil,
		nil,
	}
	return sc
}

/**
 * 创建ssh client连接
 */
func (sc *SSHClient) CreateClient() (*SSHClient, error) {
	config := &ssh.ClientConfig{
		User: sc.username,
		Auth: []ssh.AuthMethod{
			ssh.Password(sc.password),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
	}
	var err error
	sc.client, err = ssh.Dial("tcp", sc.host+":"+sc.port, config)
	util.CheckErr(err, "连接ssh client错误：")

	return sc, nil
}

/**
 * 创建ssh Session连接
 */
func (sc *SSHClient) CreatesshSession() (*SSHClient, error) {
	var err error
	sc.sshSession, err = sc.client.NewSession()
	util.CheckErr(err, "连接ssh session错误：")

	return sc, nil
}

/**
 * 创建ssh client连接
 */
func (sc *SSHClient) CreatesftpSession() (*SSHClient, error) {
	var err error
	sc.sftpSession, err = sftp.NewClient(sc.client)
	util.CheckErr(err, "连接ssh session错误：")

	return sc, nil
}

/**
 * 批量执行shell命令
 */
func (sc *SSHClient) Shell(cmds []string) string {

	stdinBuf, err := sc.sshSession.StdinPipe()
	util.CheckErr(err, "标准输入错误: ")
	stdoutBuf, err := sc.sshSession.StdoutPipe()
	util.CheckErr(err, "标准输出错误: ")

	var res bytes.Buffer
	err = sc.sshSession.Shell()
	util.CheckErr(err, "Shell状态错误: ")
	var cmd string
	for _, c := range cmds {
		if cmd != "" {
			cmd = cmd + " && " + c
		} else {
			cmd = c
		}

	}
	cmd = cmd + " && exit\n"
	// fmt.Println(cmd)
	_, err = stdinBuf.Write([]byte(cmd))
	util.CheckErr(err, cmd+"命令执行错误: ")
	reader := bufio.NewReader(stdoutBuf)
	//实时循环读取输出流中的一行内容
	for {
		line, err2 := reader.ReadString('\n')
		if err2 != nil || io.EOF == err2 {
			break
		}
		re, _ := regexp.Compile(".*(#\\s)")
		line = re.ReplaceAllString(line, "")
		fmt.Println(line)
		//输出结果存入缓冲
		res.WriteString(line)
	}
	sc.sshSession.Wait()

	return res.String()
}

// Upload 上传文件
func (sc *SSHClient) Upload(from string, to string) error {
	var uploadErr error

	srcFile, err := os.Open(from)
	util.CheckErr(err, "Upload 打开from文件错误：")
	defer srcFile.Close()

	var remoteFileName = path.Base(from)
	dstFile, err := sc.sftpSession.Create(path.Join(to, remoteFileName))
	util.CheckErr(err, "Upload 创建to文件错误：")
	defer dstFile.Close()

	buf, err := ioutil.ReadAll(srcFile)
	if err != nil {
		fmt.Println("ReadAll error : ", from)
		log.Fatal(err)

	}
	dstFile.Write(buf)

	/*buf := make([]byte, 1024)
	for {
		n, _ := srcFile.Read(buf)
		if n == 0 {
			break
		}
		dstFile.Write(buf)
	}*/
	return uploadErr
}

// Download 下载文件
func (sc *SSHClient) Download(from string, to string) error {
	var downloadErr error
	srcFile, err := sc.sftpSession.Open(from)
	if err != nil {
		log.Fatal(err)
	}
	defer srcFile.Close()

	var localFileName = path.Base(from)
	dstFile, err := os.Create(path.Join(to, localFileName))
	if err != nil {
		log.Fatal(err)
	}
	defer dstFile.Close()

	if _, err = srcFile.WriteTo(dstFile); err != nil {
		log.Fatal(err)
	}

	return downloadErr
}
